#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sun Nov  3 10:59:35 2024

@author: liangy
E-mail: liangy19@mails.tsinghua.edu.cn

For spherical lensing system

Unit System
The unit system we apply here is:
     distance         --> kiloparsec 
     angular distance --> arcsec
     mass             --> solar mass
     time             --> Gyr
     velocity         --> km/s
     
For mocking the central velocity dispersion
"""

import os 
import emcee
import corner
import numpy as np
import matplotlib.pyplot as plt
from scipy import constants
from scipy.stats import norm
from multiprocessing import Pool
from astropy.cosmology import FlatLambdaCDM
from sph_profile import sph_profile as sph
from sph_jeans import sph_jeans as sph_dyn
from sph_lensing import sph_lensing as sph_lens

"""
The physical constants required
"""
G   = 4.301e-6          # gravitational constant in (km/s)^2*kpc/M_sun
h   = 0.6774            # reduced Hubble constant
H0  = 100*h             # Hubble constant in km/s/Mpc
Om0 = 0.3089            # matter density in cosmic inventory at z = 0
c   = constants.c*1e-3  # speed of light in km/s
arcsec2rad = np.pi/180/3600    # arcsec/rad

cosmo = FlatLambdaCDM(H0,Om0)

class sph_dyn_lens_fit():
    """
    The example to fit the central velocity dispersion and Einstein radius 
    """
    def __init__(self,fit_specification):
        
        self.fit_specification = fit_specification
        
        # if the cosmology is not determined, use the Planck2016
        try:
            self.cosmo = fit_specification["cosmo"]
        except:
            self.cosmo = cosmo
        
        # basic info
        self.Re  = fit_specification["Re"] 
        self.Rap = fit_specification["Rap"]
        self.npix = fit_specification["npix"]
        self.R_bin = np.linspace(1.0/self.npix,1.0,self.npix)*self.Rap
        
        # kinematic and lensing data as well as their errors
        self.theta_ein = fit_specification["theta_ein"]
        self.vrms_ap   = fit_specification["vrms_ap"]
        self.theta_ein_err = fit_specification["theta_ein_err"]
        self.vrms_ap_err   = fit_specification["vrms_ap_err"]
        
        # the lensing arguments
        self.zd = fit_specification["zd"]
        self.zs = fit_specification["zs"]
        self.Dd = self.cosmo.angular_diameter_distance(self.zd).value*1e3
        self.Ds = self.cosmo.angular_diameter_distance(self.zs).value*1e3
        self.R_ein = self.theta_ein*arcsec2rad*self.Dd
        
        # the dynamical arguments
        self.massless_tracer = fit_specification["massless_tracer"]    
        self.polynomial  = fit_specification["polynomial"]    
        self.npoly       = fit_specification["npoly"]         
        self.anisotropy  = fit_specification["anisotropy"]
        self.beta        = fit_specification["beta"]
        self.ra          = fit_specification["ra"]
        
        # the distribution of mass and tracer
        self.tracer_model_list = fit_specification["tracer_model_list"]
        self.tracer_model_para = fit_specification["tracer_model_para"]
        self.mass_model_list   = fit_specification["mass_model_list"]
        self.tracer_model_label = fit_specification["tracer_model_label"]
        
        # mcmc arguments
        self.omega_label = fit_specification["omega_label"]
        self.omega_low_lim = fit_specification["omega_low_lim"]
        self.omega_upp_lim = fit_specification["omega_upp_lim"]
        
        # the preliminary set of lensing and dynamical modeling 
        self.lens_specification = {
                "zd":self.zd,
                "zs":self.zs,
                "cosmo":self.cosmo,
                "mass_model_list":self.mass_model_list,
                "mass_model_para":None}
        
        self.dyn_specification = {
                "massless_tracer":self.massless_tracer,
                "polynomial":self.polynomial,
                "npoly":self.npoly,
                "anisotropy":self.anisotropy,
                "tracer_model_list":self.tracer_model_list,
                "tracer_model_para":self.tracer_model_para,
                "tracer_model_label":self.tracer_model_label,
                "mass_model_list":self.mass_model_list,
                "mass_model_para":None,
                "beta":self.beta,
                "ra":self.ra}
        
        self.sph_ein_2d = sph(self.R_ein,dim = "2D")
        self.sph_ein_3d = sph(self.R_ein,dim = "3D")
        self.sph_5kpc_2d = sph(5.0,dim = "2D")
        self.sph_5kpc_3d = sph(5.0,dim = "3D")
        
    def posterior_sampling(self,nwalkers=20,nsamples=3000,nwarmup=500,ndim=2, 
                           progress = False):
        self.nwalkers = nwalkers
        self.nsamples = nsamples
        self.nwarmup  = nwarmup
        self.ndim = ndim
        
        # generate the initial state of the walkers
        self.omega0  = [self.random_omega() for _ in range(nwalkers)]
        
        # sample in the parameter space with parallelization
        with Pool(processes = nwalkers) as pool:
            self.sampler = emcee.EnsembleSampler(
                    self.nwalkers,self.ndim,
                    self.ln_posterior,pool = pool)
            self.sampler.run_mcmc(self.omega0,nsamples,progress = progress)
            
        # the markov chains under our posterior
        self.warm_chain = self.sampler.chain[:,self.nwarmup:,:].reshape((-1,self.ndim))
        
    def posterior_estimating(self,infer = 'median'):
        """
        attention: the shape of the np.ndarray sampler.get_chain() and sampler.chain is not the same
        in sampler.chain, [nwalkers,nsamples,ndim]
        in sampler.get_chain(), [nsamples,nwalkers,ndim]
        in sampler.get_log_prob(), [nsamples,nwalkers]
        the method reshape(-1,ndim) is chain by chain in sampler.chain
        the method reshape(-1,ndim) is step by step in sampler.get_chain()
        """
        if infer == 'median':
            self.omega_middle = np.percentile(self.warm_chain,50,axis = 0)
            self.omega_lower  = np.percentile(self.warm_chain,16,axis = 0)
            self.omega_upper  = np.percentile(self.warm_chain,84,axis = 0)
            return self.omega_middle,self.omega_lower,self.omega_upper
        if infer == 'max_likelihood':
            max_likelihood = self.sampler.get_log_prob().argmax()
            self.omega_middle = self.sampler.get_chain().reshape(-1,self.ndim)[max_likelihood]
            self.omega_lower  = np.percentile(self.warm_chain,16,axis = 0)
            self.omega_upper  = np.percentile(self.warm_chain,84,axis = 0)
            return self.omega_middle,self.omega_lower,self.omega_upper
        if infer == 'average':
            self.omega_middle = np.average(self.warm_chain,axis = 0)
            self.omega_lower  = self.omega_middle - np.std(self.warm_chain,axis = 0)
            self.omega_upper  = self.omega_middle + np.std(self.warm_chain,axis = 0)
            return self.omega_middle,self.omega_lower,self.omega_upper
    
    def ln_posterior(self,omega):
        if all(self.omega_low_lim<omega) and all(omega<self.omega_upp_lim):
            ln_post = self.ln_likelihood(omega) 
            return ln_post
        else:
            return -np.inf
    
    def ln_likelihood(self,omega):
        if self.massless_tracer:
            lg_Mein,eta = omega
            
            # convert mcmc parameters as sph_profile parameterization
            rho_tot_ein = 10**lg_Mein/self.sph_ein_2d.M_power_law(1.0,self.R_ein,eta)
            mass_model_para = [[rho_tot_ein,self.R_ein,eta]]
            
            # re-normlization of accelerated_Mnu_roots_bundle
            Mnu = 1.0
        
        else:
            lg_Mst,lg_Mdm5 = omega
            
            # convert mcmc parameters as sph_profile parameterization
            M_st = 10**lg_Mst
            rho_dm_0 = 10**lg_Mdm5/self.sph_5kpc_2d.M_two_power_law(1.0,10*self.Re,1.0,3.0)
            mass_model_para = [[M_st,self.Re,4.0],[rho_dm_0,10*self.Re,1.0,3.0]]
            
            # re-normalization of accelerated_Mnu_roots_bundle
            Mnu = M_st
            
        # updata the modeling parameters at each step
        lens_specification = self.lens_specification
        dyn_specification  = self.dyn_specification
        lens_specification["mass_model_para"] = mass_model_para
        dyn_specification["mass_model_para"]  = mass_model_para
        
        # lensing modeling
        lensing  = sph_lens(lens_specification)
        theta_ein_mod = lensing.get_einstein_radius(self.theta_ein)
        
        # dynamical modeling
        vrms  = np.zeros_like(self.R_bin)
        Sigma = np.zeros_like(self.R_bin)
        for i in range(self.npix):
            dynamics_i = sph_dyn(self.R_bin[i],dyn_specification)
            vrms[i],Sigma[i] = dynamics_i.accelerated_solver(
                    self.accelerated_nu_roots_bundle[i],
                    self.accelerated_Mnu_roots_bundle[i]*Mnu)
        vrms_ap_mod = np.average(vrms,weights = self.R_bin*Sigma)
        
        # gaussian likelihood
        ln_gauss_lens = norm.logpdf(
                self.theta_ein,
                loc = theta_ein_mod,
                scale = self.theta_ein_err)
        ln_gauss_dyn  = norm.logpdf(
                self.vrms_ap,  
                loc = vrms_ap_mod,  
                scale = self.vrms_ap_err)
        return ln_gauss_lens+ln_gauss_dyn
        
    def acceleration_mcmc(self):
        """
        attention: to keep the accelerated_Mnu_roots_bundle is model-independent,
            it only obtained with tracer_para (normalization factor is 1.0).
            please re-normalize accelerated_Mnu_roots_bundle during mcmc sampling.
        """
        if self.massless_tracer:
            accelerated_nu_roots_bundle  = np.zeros((self.npix,self.npoly),dtype = float)
            accelerated_Mnu_roots_bundle = np.zeros((self.npix,self.npoly),dtype = float)
            # convert mcmc parameters as sph_profile parameterization
            
            # updata the modeling parameters at each step
            dyn_specification  = self.dyn_specification
            
            # generate the tracer density and mass at the points corresponding to polunomial roots 
            for i in range(self.npix):
                dynamics_i = sph_dyn(self.R_bin[i],dyn_specification)
                accelerated_nu_roots_bundle[i] = dynamics_i.acceleration()
            self.accelerated_nu_roots_bundle = accelerated_nu_roots_bundle
            self.accelerated_Mnu_roots_bundle = accelerated_Mnu_roots_bundle
        
        else:
            accelerated_nu_roots_bundle  = np.zeros((self.npix,self.npoly),dtype = float)
            accelerated_Mnu_roots_bundle = np.zeros((self.npix,self.npoly),dtype = float)
            
            # updata the modeling parameters at each step
            dyn_specification  = self.dyn_specification
            
            
            # generate the tracer density and mass at the points corresponding to polunomial roots 
            for i in range(self.npix):
                dynamics_i = sph_dyn(self.R_bin[i],dyn_specification)
                accelerated_nu_roots_bundle[i],accelerated_Mnu_roots_bundle[i] = dynamics_i.acceleration()
            self.accelerated_nu_roots_bundle  = accelerated_nu_roots_bundle
            self.accelerated_Mnu_roots_bundle = accelerated_Mnu_roots_bundle
        
    def random_omega(self):
        return np.random.uniform(low  = self.omega_low_lim,
                                 high = self.omega_upp_lim,
                                 size = self.ndim)
    
    def auto_correlation(self):
        tau = self.sampler.get_autocorr_time(quiet = True)
        return tau
    
if __name__ == "__main__":
    """
    An example of this lensing-dynamics modeling routine
    """
    
    """
    The basic setup of the matplotlib
    """
    # switch the matplotlib backend to fit the environment of computing cluster
    plt.switch_backend("agg")
    font_scale = 2.0
    plt.rcParams["font.family"]      = "serif"
    plt.rcParams["font.serif"]       = "Times New Roman"
    plt.rcParams["mathtext.rm"]      = "serif"
    plt.rcParams["mathtext.it"]      = "serif:italic"
    plt.rcParams["mathtext.bf"]      = "serif:bold"
    plt.rcParams["mathtext.fontset"] = "custom"
    plt.rcParams["xtick.direction"]  = "in"
    plt.rcParams["ytick.direction"]  = "in"
    plt.rcParams["xtick.labelsize"]  = 6.0*font_scale
    plt.rcParams["ytick.labelsize"]  = 6.0*font_scale
    plt.rcParams["axes.linewidth"]   = 0.5*font_scale
    plt.rcParams["axes.labelsize"]   = 9.0*font_scale
    plt.rcParams["axes.titlesize"]   = 5.0*font_scale
    plt.rcParams["font.size"]        = 9.0*font_scale
    plt.rcParams["legend.fontsize"]  = 5.0*font_scale
    plt.rcParams["xtick.top"]    = True
    plt.rcParams["xtick.bottom"] = True 
    plt.rcParams["ytick.left"]   = True 
    plt.rcParams["ytick.right"]  = True
    plt.rcParams["xtick.minor.top"]    = True
    plt.rcParams["xtick.minor.bottom"] = True 
    plt.rcParams["ytick.minor.left"]   = True 
    plt.rcParams["ytick.minor.right"]  = True
    plt.rcParams["xtick.minor.visible"] = True
    plt.rcParams["ytick.minor.visible"] = True
    plt.rcParams["xtick.major.size"] = 3.0*font_scale
    plt.rcParams["ytick.major.size"] = 3.0*font_scale
    plt.rcParams["xtick.minor.size"] = 1.5*font_scale
    plt.rcParams["ytick.minor.size"] = 1.5*font_scale
    plt.rcParams["xtick.major.width"] = 0.5*font_scale
    plt.rcParams["ytick.major.width"] = 0.5*font_scale
    plt.rcParams["xtick.minor.width"] = 0.35*font_scale
    plt.rcParams["ytick.minor.width"] = 0.35*font_scale
    
    
    current_dir = os.path.dirname(os.path.abspath(__file__))
    Galaxy_ID = "SL2SJ021411−040502"
    cosmo_SL2S = FlatLambdaCDM(70.0,0.3)
    
    # scenario 1: tracer massless, 
    # signle mass component: power law
    fit_specification_Mein_eta = {
            "cosmo":cosmo_SL2S,
            "Re": 6.290,    # in kpc
            "Rap":3.145,    # in kpc
            "theta_ein":1.410,    # in arcsec
            "theta_ein_err":0.04, # in arcsec
            "vrms_ap":238.0,      # in km/s
            "vrms_ap_err":15.0,    # in km/s
            "zd":0.609,
            "zs":1.88,
            "npix":20,
            "npoly":30,
            "massless_tracer":True,
            "polynomial":"legendre",
            "anisotropy":"const",
            "beta":0.0,
            "ra":None,
            "tracer_model_list":["sersic"],    # deV
            "tracer_model_para":[[1.0,6.290,4.0]],
            "mass_model_list":["power_law"],    # deV/PL
            "tracer_model_label":[False],
            "omega_label":[r"$\log{M_{\rm E}/\rm M_\odot}$",r"$\eta$"],
            "omega_low_lim":[10.0,1.2],
            "omega_upp_lim":[12.5,2.8]}
    
    lens_dyn_Mein_eta = sph_dyn_lens_fit(fit_specification_Mein_eta)
    lens_dyn_Mein_eta.acceleration_mcmc()
    lens_dyn_Mein_eta.posterior_sampling()
    lens_dyn_Mein_eta.posterior_estimating()
    
    corner_Mein_eta = corner.corner(
            lens_dyn_Mein_eta.warm_chain,
            bins = 30,quantiles=[0.16,0.5,0.84],
            color = 'seagreen',smooth = True,smooth1d = True,plot_contours = True,
            labels = lens_dyn_Mein_eta.omega_label,use_math_text = True,
            show_titles=True)
    corner_Mein_eta.savefig(current_dir+'/{}_Mein_eta_corner.pdf'.format(Galaxy_ID),bbox_inches = 'tight')
    
    # scenario 2: tracer contributes mass
    # two mass components: deV(stars/tracer), fixed NFW(dark matter) 
    fit_specification_Mst_Mdm5 = {
            "cosmo":cosmo_SL2S,
            "Re": 6.290,    # in kpc
            "Rap":3.145,    # in kpc
            "theta_ein":1.410,    # in arcsec
            "theta_ein_err":0.04, # in arcsec
            "vrms_ap":238.0,      # in km/s
            "vrms_ap_err":15.0,    # in km/s
            "zd":0.609,
            "zs":1.88,
            "npix":20,
            "npoly":30,
            "massless_tracer":False,
            "polynomial":"legendre",
            "anisotropy":"const",
            "beta":0.0,
            "ra":None,
            "tracer_model_list":["sersic"],    # deV
            "tracer_model_para":[[1.0,6.290,4.0]],
            "mass_model_list":["sersic","two_power_law"],    # deV+NFW
            "tracer_model_label":[True,False],
            "omega_label":[r"$\log{M_\star/\rm M_\odot}$",r"$\log{M_{\rm DM5}/\rm M_\odot}$"],
            "omega_low_lim":[10.0,10.0],
            "omega_upp_lim":[12.5,12.5]}
    
    lens_dyn_Mst_Mdm5 = sph_dyn_lens_fit(fit_specification_Mst_Mdm5)
    lens_dyn_Mst_Mdm5.acceleration_mcmc()
    lens_dyn_Mst_Mdm5.posterior_sampling()
    lens_dyn_Mst_Mdm5.posterior_estimating()
    
    corner_Mst_Mdm5 = corner.corner(
            lens_dyn_Mst_Mdm5.warm_chain,
            bins = 30,quantiles=[0.16,0.5,0.84],
            color = 'seagreen',smooth = True,smooth1d = True,plot_contours = True,
            labels = lens_dyn_Mst_Mdm5.omega_label,use_math_text = True,
            show_titles=True)
    corner_Mst_Mdm5.savefig(current_dir+'/{}_Mst_Mdm5_corner.pdf'.format(Galaxy_ID),bbox_inches = 'tight')
    